import torch 
from models import MultiLayerPerceptron
import numpy as np
import torch.nn.functional as F

# 暂时不知道放在哪个初始化下合适
learning_rate = 0.01
gamma = 0.98
epsilon = 0.01
target_update = 10

class DQN():
    # 编写初始化
    def __init__(self, num_states, num_actions, config):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        self.action_dim = num_actions
        self.state_dim = num_states
        # 网络结构设计
        self.eval_net = MultiLayerPerceptron(config).to(self.device)
        self.target_eval_net = MultiLayerPerceptron(config).to(self.device)
        
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(),
                                          lr=learning_rate)
        self.gamma = gamma  # 折扣因子
        self.epsilon = epsilon  # epsilon-贪婪策略
        self.target_update = target_update  # 目标网络更新频率
        self.count = 0  # 计数器,记录更新次数
    
        
    #编写选择动作
    def choose_action(self, state):
        if np.random.random() < self.epsilon:
            action = np.random.randint(self.action_dim)
        else:
            state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
            action = self.eval_net(state).argmax().item()
        return action
        
    #编写agent算法训练过程
    def learn(self, transition_dict):
        # sars d
        states = torch.tensor(transition_dict['states'],dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'],dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'],dtype=torch.float).view(-1, 1).to(self.device)
        # 计算dqn_loss
        q_values = self.eval_net(states).gather(1, actions)  # Q值
        max_next_q_values = self.target_eval_net(next_states).max(1)[0].view(-1, 1) # 下个状态的最大Q值
        q_targets = rewards + self.gamma * max_next_q_values * (1 - dones)  # TD误差目标
        dqn_loss = torch.mean(F.mse_loss(q_values, q_targets))  # 均方误差损失函数
        # 更新网络参数
        self.optimizer.zero_grad()  # PyTorch中默认梯度会累积,这里需要显式将梯度置为0
        dqn_loss.backward()  # 反向传播更新参数
        self.optimizer.step()

        if self.count % self.target_update == 0:
            self.target_eval_net.load_state_dict(
                self.eval_net.state_dict())  # 更新目标网络
        self.count += 1
       

    def save(self, directory='results\dqn'):
        torch.save(self.eval_net.state_dict(), directory + 'dqn.pth')
        # print("====================================")
        # print("Model has been saved...")
        # print("====================================")

    def load(self, directory, i):
        self.eval_net.load_state_dict(torch.load(directory + 'dqn{}.pth'.format(i)))
        print("====================================")
        print("Model has been loaded...")
        print("====================================")